import os
import sys
import json
import time
from dotenv import load_dotenv
from loguru import logger

from src.utils import setup_logger
from src.mcp import MCPClient

# 加载环境变量
load_dotenv()

# 设置日志
setup_logger()

def print_banner():
    """打印欢迎横幅"""
    banner = """
================================================================================
                         智能SQL数据库助手客户端
================================================================================
    使用自然语言查询和操作您的数据库，由DeepSeek大模型提供支持。
    
    命令:
    - help                 显示帮助信息
    - connect              连接到服务器
    - disconnect           断开服务器连接
    - info                 获取数据库信息
    - query <自然语言查询>   执行自然语言查询
    - sql <SQL查询>         执行SQL查询
    - exit/quit            退出客户端
================================================================================
"""
    print(banner)

def print_help():
    """打印帮助信息"""
    help_text = """
可用命令:
    help                 - 显示此帮助信息
    connect              - 连接到服务器
    disconnect           - 断开服务器连接
    info                 - 获取数据库结构信息
    query <自然语言查询>   - 使用自然语言查询数据库
      选项:
        --format=<格式>   - 指定输出格式 (text, json, html, csv)
        --no-explain     - 不显示SQL解释
        --no-analysis    - 不显示结果分析
        --no-viz         - 不显示可视化建议
    
    sql <SQL查询>         - 直接执行SQL查询
      选项:
        --format=<格式>   - 指定输出格式 (text, json, html, csv)
        --no-validate    - 不验证SQL
    
    clear                - 清屏
    exit, quit           - 退出客户端

示例:
    query 显示销售表中最近10条记录
    query 找出订单量最多的前5个客户 --format=json
    sql SELECT * FROM customers LIMIT 5
"""
    print(help_text)

def parse_command(command_line):
    """解析命令行"""
    parts = command_line.strip().split(maxsplit=1)
    command = parts[0].lower() if parts else ""
    
    args = ""
    options = {}
    
    if len(parts) > 1:
        # 分离参数和选项
        remaining = parts[1]
        
        # 提取选项
        option_parts = []
        for part in remaining.split():
            if part.startswith("--"):
                option_part = part[2:]  # 移除前缀 '--'
                if "=" in option_part:
                    key, value = option_part.split("=", 1)
                    options[key] = value
                else:
                    options[option_part] = True
            else:
                option_parts.append(part)
        
        args = " ".join(option_parts)
    
    return command, args, options

def format_response(response, command):
    """格式化并打印响应"""
    # 检查响应结构
    if not isinstance(response, dict):
        print(f"\n错误: 接收到无效的响应格式: {response}")
        return
    
    # 检查错误信息
    if "error" in response:
        print(f"\n错误: {response['error']}")
        return
    
    # 检查状态字段
    if "status" in response and response["status"] == "error":
        print(f"\n错误: {response.get('message', '未知错误')}")
        return
    
    # 对于info命令，尝试从响应中获取data字段
    if command == "info":
        # 如果响应中有data字段，使用data中的内容
        if "data" in response:
            response = response["data"]
        
        # 检查必要的元数据字段
        if "metadata" not in response:
            print("\n错误: 服务器返回的数据中缺少元数据信息")
            print(f"返回的数据: {json.dumps(response, indent=2, ensure_ascii=False)[:200]}...")
            return
            
        print("\n=== 数据库信息 ===")
        print(f"表数量: {len(response['metadata'])}")
        
        for table_name, table_info in response["metadata"].items():
            row_count = response.get("table_counts", {}).get(table_name, "未知")
            print(f"\n表: {table_name} (约 {row_count} 行)")
            
            print("  列:")
            for column in table_info["schema"]:
                nullable = "" if column["nullable"] else " NOT NULL"
                default = f" DEFAULT {column['default']}" if column.get("default") else ""
                print(f"    - {column['name']} ({column['type']}){nullable}{default}")
            
            if table_info["relationships"]:
                print("  外键关系:")
                for rel in table_info["relationships"]:
                    print(f"    - {','.join(rel['constrained_columns'])} -> {rel['referred_table']}.{','.join(rel['referred_columns'])}")
            
            sample_data = response.get("sample_data", {}).get(table_name, [])
            if sample_data:
                print(f"  样本数据 ({len(sample_data)} 行):")
                for i, row in enumerate(sample_data):
                    print(f"    行 {i+1}: {json.dumps(row, ensure_ascii=False)}")
    
    # 对于query命令，也需要检查data字段
    elif command == "query":
        if "data" in response:
            response = response["data"]
            
        # 检查必要字段
        if "query" not in response or "generated_sql" not in response:
            print("\n错误: 服务器返回的查询结果格式不正确")
            print(f"返回的数据: {json.dumps(response, indent=2, ensure_ascii=False)[:200]}...")
            return
            
        print("\n=== 查询结果 ===")
        print(f"自然语言查询: {response['query']}")
        print(f"\n生成的SQL: {response['generated_sql']}")
        
        if "explanation" in response:
            print(f"\nSQL解释:\n{response['explanation']}")
        
        print(f"\n结果:\n{response.get('formatted_results', '无结果')}")
        
        if "analysis" in response and isinstance(response["analysis"], dict) and "analysis" in response["analysis"]:
            print(f"\n数据分析:\n{response['analysis']['analysis']}")
        elif "analysis" in response and isinstance(response["analysis"], str):
            print(f"\n数据分析:\n{response['analysis']}")
        
        if "visualization" in response:
            if isinstance(response["visualization"], dict) and "visualization_suggestion" in response["visualization"]:
                print(f"\n可视化建议:\n{response['visualization']['visualization_suggestion']}")
            elif isinstance(response["visualization"], str):
                print(f"\n可视化建议:\n{response['visualization']}")
    
    # 对于sql命令，也需要检查data字段
    elif command == "sql":
        if "data" in response:
            response = response["data"]
            
        # 检查必要字段
        if "sql" not in response:
            print("\n错误: 服务器返回的SQL结果格式不正确")
            print(f"返回的数据: {json.dumps(response, indent=2, ensure_ascii=False)[:200]}...")
            return
            
        print("\n=== SQL执行结果 ===")
        print(f"SQL查询: {response['sql']}")
        
        if "validation" in response and response["validation"]:
            if response["validation"].get("warnings", []):
                print("\n警告:")
                for warning in response["validation"]["warnings"]:
                    print(f"  - {warning}")
        
        print(f"\n结果:\n{response.get('formatted_results', '无结果')}")
    
    # 对于其他命令，直接打印数据
    else:
        if "data" in response:
            response = response["data"]
        print(f"\n响应数据:\n{json.dumps(response, indent=2, ensure_ascii=False)}")

def main():
    """主程序"""
    print_banner()
    
    # 创建MCP客户端
    client = MCPClient()
    connected = False
    
    while True:
        try:
            # 提示符
            status = "已连接" if connected else "未连接"
            prompt = f"\n智能SQL助手 [{status}]> "
            
            # 获取用户输入
            user_input = input(prompt)
            if not user_input.strip():
                continue
            
            # 解析命令
            command, args, options = parse_command(user_input)
            
            # 处理命令
            if command in ["exit", "quit"]:
                print("正在退出客户端...")
                if connected:
                    client.disconnect()
                break
            
            elif command == "help":
                print_help()
            
            elif command == "clear":
                # 清屏
                os.system('cls' if os.name == 'nt' else 'clear')
                print_banner()
            
            elif command == "connect":
                if connected:
                    print("已经连接到服务器")
                else:
                    # 重新加载环境变量，确保使用最新配置
                    from dotenv import load_dotenv
                    load_dotenv(override=True)
                    
                    host = os.getenv("MCP_HOST", "localhost")
                    port = int(os.getenv("MCP_PORT", "8080"))
                    print(f"正在连接到服务器 {host}:{port}...")
                    
                    # 创建新的客户端连接，确保使用最新参数
                    client = MCPClient(host=host, port=port)
                    
                    if client.connect():
                        connected = True
                        print(f"已成功连接到服务器 {host}:{port}")
                    else:
                        print("连接服务器失败")
            
            elif command == "disconnect":
                if not connected:
                    print("未连接到服务器")
                else:
                    client.disconnect()
                    connected = False
                    print("已断开与服务器的连接")
            
            elif command == "info":
                if not connected:
                    print("请先连接到服务器")
                    continue
                
                print("正在获取数据库信息...")
                response = client.get_database_info()
                format_response(response, command)
            
            elif command == "query":
                if not connected:
                    print("请先连接到服务器")
                    continue
                
                if not args:
                    print("请提供自然语言查询")
                    continue
                
                # 准备请求
                request = {
                    "action": "query",
                    "query": args,
                    "format": options.get("format", "text"),
                    "include_explanation": "no-explain" not in options,
                    "include_analysis": "no-analysis" not in options,
                    "include_visualization": "no-viz" not in options
                }
                
                print(f"正在处理查询: {args}")
                response = client.send_request(request)
                format_response(response, command)
            
            elif command == "sql":
                if not connected:
                    print("请先连接到服务器")
                    continue
                
                if not args:
                    print("请提供SQL查询")
                    continue
                
                # 准备请求
                request = {
                    "action": "sql",
                    "sql": args,
                    "format": options.get("format", "text"),
                    "validate": "no-validate" not in options
                }
                
                print(f"正在执行SQL: {args}")
                response = client.send_request(request)
                format_response(response, command)
            
            else:
                print(f"未知命令: {command}")
                print("输入 'help' 查看可用命令")
        
        except KeyboardInterrupt:
            print("\n操作已取消")
        except Exception as e:
            print(f"\n发生错误: {str(e)}")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n程序已终止")
    except Exception as e:
        logger.error(f"程序出现未处理的异常: {str(e)}")
        sys.exit(1) 